#include "rose.h"
#include <sstream>
#include <string>
#include <iostream>
#include "Argument.h"
#include "FunctionCall.h"
#include "FunctionContext.h"
#include "AnalyzeTraversal.h"
#include "Reference.h"
#include "VariableWrapper.h"
#include "ConfirmFunctions.h"
#include "CollectRefs.h"

#define DO_PRINT

using namespace std;

void analyzeTraversal::init(FunctionContext **fc, int N) {
    current_function = NULL;
    current_function_index = -1;

    all_functions = fc;
    n_functions = N;

    scope_depth = 0;
    scope_stack = NULL;

    globals = NULL;
    n_globals = 0;

    structs = NULL;
    n_structs = 0;

    current_struct = NULL;
    current_struct_index = -1;

}

analyzeTraversal::analyzeTraversal(FunctionContext **set_functions, int init_num_functions) {
    init(set_functions, init_num_functions);
}

analyzeTraversal::analyzeTraversal()  {
    init(NULL, NULL);
}

void analyzeTraversal::print() {
    for(int i = 0; i < n_functions ; i++) {
        if(all_functions[i]->isConfirmed() == 1) {
            string str = all_functions[i]->str();
            cout<<str<<endl;
        }
    }
    cout<<endl<<"We have "<<n_globals<<" global(s): ";
    for(int i = 0; i < n_globals; i++) {
        cout<<globals[i]->str()<<" ";
    }
    cout<<endl;
    cout<<"We have "<<n_structs<<" struct definition(s): "<<endl;
    for(int i = 0; i < n_structs; i++) {
        cout<<structs[i]->str()<<endl;
    }
    cout<<endl;
}

SgDeclarationStatement* analyzeTraversal::resolveSymbol(SgVarRefExp *ref) {
    SgSymbol *sym = ref->get_symbol();
    int current_depth = scope_depth-1;
    SgScopeStatement *scope = scope_stack[current_depth];
    while(current_depth >= 0 && scope->lookup_variable_symbol(sym->get_name()) == NULL) {
        current_depth--;
        scope = scope_stack[current_depth];
    }
    if(current_depth < 0) return NULL;
    SgVariableSymbol *var_sym = scope->lookup_variable_symbol(sym->get_name());
    SgInitializedName *init_name = var_sym->get_declaration();
    SgDeclarationStatement *decl = init_name->get_declaration();
    if(decl == NULL) {
        // find in globals
        for(int i = 0; i < n_globals; i++) {
            VariableWrapper *global = globals[i];
            if(global->str() == ref->get_symbol()->get_name().str()) {
                return (SgDeclarationStatement *)(globals[i]->get_decl());
            }
        }
    }
    return decl;
}

void analyzeTraversal::appendReads(Reference **refs, int n) {
    if(n > 0) {
        current_function->appendReads(refs, n);
    }
}

void analyzeTraversal::addStruct(SgClassDefinition *definition) {
    current_struct = new StructWrapper(definition);
    current_struct_index = n_structs;

    structs = (StructWrapper **)realloc(structs, (n_structs + 1) * sizeof(StructWrapper *));
    structs[n_structs] = current_struct;
    n_structs = n_structs + 1;
}

void analyzeTraversal::addGlobal(SgVariableDeclaration *global) {

    SgInitializedNamePtrList& vars = global->get_variables();
    for(unsigned int i = 0; i < vars.size(); i++) {
        SgInitializedName *v = vars[i];
        globals = (VariableWrapper **)realloc(globals, (n_globals+1) * sizeof(VariableWrapper *));
        globals[n_globals] = new VariableWrapper(v, global, NULL);
        n_globals++;
    }
}

int analyzeTraversal::addFunction(SgFunctionDeclaration *decl) {
    if(findFunction(decl->get_name().str()) != -1) return findFunction(decl->get_name().str());
    all_functions = (FunctionContext **)realloc(all_functions, (n_functions + 1) * sizeof(FunctionContext *));
    all_functions[n_functions] = new FunctionContext(decl, this);
    n_functions = n_functions + 1;
    return n_functions-1;
}

VariableWrapper *analyzeTraversal::addLocalVar(SgVariableDeclaration *decl) {
    if(current_function != NULL) {
        return current_function->addLocal(decl);
    } else if(current_struct != NULL ) {
        current_struct->addField(decl);
    } else {
        addGlobal(decl);
    }
    return NULL;
}

void analyzeTraversal::appendWrites(Reference **refs, int n) {
    if(current_function != NULL) {
        current_function->appendWrites(refs, n);
    }
}

int analyzeTraversal::findFunction(FullFunctionCall *call) {
    for(int i = 0; i < n_functions; i++) {
        if(all_functions[i]->get_name() == call->get_name()) {
            return i;
        }
    }
    return -1;
}

int analyzeTraversal::findFunction(SgFunctionDefinition *def) {
    for(int i = 0; i < n_functions; i++) {
        if(all_functions[i]->get_decl() == def->get_declaration()) {
            return i;
        }
    }
    return -1;
}

int analyzeTraversal::findFunction(SgFunctionDeclaration *decl) {
    for(int i = 0; i < n_functions; i++) {
        if(all_functions[i]->get_decl() == decl) {
            return i;
        }
    }
    return -1;
}

int analyzeTraversal::findFunction(string name) {
    for(int i = 0; i < n_functions; i++) {
        if(all_functions[i]->get_decl()->get_name().str() == name) {
            return i;
        }
    }
    return -1;
}

Reference *analyzeTraversal::generateReference(SgExpression *arg, SgVarRefExp **var, SgDeclarationStatement **val_passed, VariableWrapper **var_passed) {

    *var = isSgVarRefExp(arg);
    SgPointerDerefExp *deref = isSgPointerDerefExp(arg);
    SgDotExp *dot = isSgDotExp(arg);
    SgArrowExp *arrow = isSgArrowExp(arg);

    Reference *head = NULL;
    Reference *base = NULL;
    if(deref != NULL || dot != NULL || arrow != NULL) {
        while(isSgDotExp(arg) != NULL || isSgPointerDerefExp(arg) != NULL || isSgArrowExp(arg)) {
            if(isSgDotExp(arg) != NULL) {
                // struct
                SgVarRefExp *field_ref = isSgVarRefExp(isSgDotExp(arg)->get_rhs_operand());
                if(head == NULL) {
                    head = new Reference(STRUCT);
                    head->setField(field_ref->get_symbol()->get_name().str());
                    base = head;
                } else {
                    Reference *tmp = new Reference(STRUCT);
                    tmp->setField(field_ref->get_symbol()->get_name().str());
                    base->setChild(tmp);
                    base = tmp;
                }
                arg = isSgDotExp(arg)->get_lhs_operand();
            } else if(isSgArrowExp(arg) != NULL) {
                SgVarRefExp *field_ref = isSgVarRefExp(arrow->get_rhs_operand());
                if(head == NULL) {
                    head = new Reference(STRUCT);
                    head->setField(field_ref->get_symbol()->get_name().str());
                    base = head;
                    Reference *tmp = new Reference(POINTER);
                    base->setChild(tmp);
                    base = tmp;
                } else {
                    Reference *tmp = new Reference(STRUCT);
                    tmp->setField(field_ref->get_symbol()->get_name().str());
                    base->setChild(tmp);
                    base = tmp;
                    tmp = new Reference(POINTER);
                    base->setChild(tmp);
                    base = tmp;
                }
                arg = arrow->get_lhs_operand();
            } else {
                // pointer
                if(head == NULL) {
                    head = new Reference(POINTER);
                    base = head;
                } else {
                    Reference *tmp = new Reference(POINTER);
                    base->setChild(tmp);
                    base = tmp;
                }
                arg = isSgPointerDerefExp(arg)->get_operand();

            }
        }
        *var = isSgVarRefExp(arg);
    }
    if(base == NULL) {
        head = base = new Reference(VARIABLE);
    } else {
        Reference *tmp = new Reference(VARIABLE);
        base->setChild(tmp);
        base = tmp;
    }

    if(*var != NULL) {
        SgDeclarationStatement *decl = resolveSymbol(*var);
        if(decl != NULL) { 
            *val_passed = decl;
            SgFunctionParameterList *params = isSgFunctionParameterList(decl);
            if(params != NULL) {
                string name = (*var)->get_symbol()->get_name().str();
                VariableWrapper *new_var = current_function->findVarByName(name);
                if(var != NULL) {
                    *var_passed = new_var;
                }
            } else {
                VariableWrapper *new_var = current_function->findVarByDecl(decl);
                if(new_var != NULL) {
                    *var_passed = new_var;
                } 
            }
        }
    }
    if(*var_passed != NULL) {
        head->setVar(*var_passed);
    } else head = NULL;

    return head;
}


void analyzeTraversal::preOrderVisit(SgNode *node ) {
    SgFunctionDeclaration *func_decl = isSgFunctionDeclaration(node);
    SgFunctionDefinition *func_def = isSgFunctionDefinition(node);
    SgFunctionCallExp *call = isSgFunctionCallExp(node);
    SgType *type = isSgType(node);
    SgVariableDeclaration *var_decl = isSgVariableDeclaration(node);
    SgVariableDefinition *var_def = isSgVariableDefinition(node);
    SgScopeStatement *scope = isSgScopeStatement(node);
    SgVarRefExp *var_ref = isSgVarRefExp(node);
    SgClassDefinition *struct_def = isSgClassDefinition(node);
    SgAssignOp *assign = isSgAssignOp(node);
    SgReturnStmt *return_stmt = isSgReturnStmt(node);
    SgScopeStatement *current_scope = NULL;

    if(scope_depth != 0) {
        current_scope = scope_stack[scope_depth-1];
    }

    if(func_decl != NULL ) {
#ifdef DO_PRINT
        cout<<"Function Declaration..."<<endl;
#endif
        addFunction(func_decl);

    } else if(func_def != NULL) {
#ifdef DO_PRINT
        cout<<"Function Definition..."<<endl;
#endif
        current_function_index = findFunction(func_def->get_declaration()->get_name().str());

        current_function = all_functions[current_function_index];
        current_function->confirmFunction();
        current_scope = func_def;
    }else if(call != NULL) {
#ifdef DO_PRINT
        cout<<"Function call..."<<endl;
#endif
        if(current_function_index != -1) {
            SgExprListExp *args = call->get_args();
            SgExpressionPtrList& exprs = args->get_expressions();
            SgDeclarationStatement **vals_passed = (SgDeclarationStatement **)malloc(sizeof(SgDeclarationStatement *) * exprs.size());
            VariableWrapper **vars_passed = (VariableWrapper **)malloc(sizeof(VariableWrapper *) * exprs.size());
            Reference **refs = (Reference **)malloc(sizeof(Reference *) * exprs.size());
            memset(vals_passed, 0x00, sizeof(SgDeclarationStatement *) * exprs.size());
            memset(vars_passed, 0x00, sizeof(VariableWrapper *) * exprs.size());
            memset(refs, 0x00, sizeof(Reference *) * exprs.size());

            for(unsigned int k = 0; k < exprs.size(); k++) {
                SgExpression *arg = exprs[k];
                SgVarRefExp *var = isSgVarRefExp(arg);
                refs[k] = generateReference(arg, &var, vals_passed+k, vars_passed+k);
                //cout<<"HAVE "<<arg->variantT()<<" in "<<current_function->get_name()<<":"<<call->getAssociatedFunctionDeclaration()->get_name().str()<<" "<<n_derefs<<endl;
/*
                if(var != NULL) {
                    SgDeclarationStatement *decl = resolveSymbol(var);
                    if(decl == NULL) { 
                        cout<<"DECL IS NULL "<<var->get_symbol()->get_name()<<endl;
                    } else {
                        vals_passed[k] = decl;
                        SgFunctionParameterList *params = isSgFunctionParameterList(decl);
                        if(params != NULL) {
                            string name = var->get_symbol()->get_name().str();
                            VariableWrapper *var = current_function->findVarByName(name);
                            if(var != NULL) {
                                vars_passed[k] = var;
                            }
                        } else {
                            VariableWrapper *var = current_function->findVarByDecl(decl);
                            if(var != NULL) {
                                vars_passed[k] = var;
                            } else {
                                cout<<"VAR IS NULL "<<endl;
                            }
                        }
                    }
                }
                if(vars_passed[k] != NULL) {
                    refs[k]->setVar(vars_passed[k]);
                } else refs[k] = NULL;
                */
            }
            FullFunctionCall *fullCall = new FullFunctionCall(call, vals_passed, vars_passed, refs, exprs.size());
            current_function->addCall(fullCall);
            FunctionContext *called_function = all_functions[findFunction(fullCall)];
            if(called_function->isConfirmed()) {
                called_function->addCaller(current_function, fullCall);
            }
        }
    } else if(type != NULL) {
#ifdef DO_PRINT
        cout<<"Visiting a type in regular traversal..."<<endl;
#endif
    } else if(var_decl != NULL) {
#ifdef DO_PRINT
        cout<<"Variable declaration..."<<endl;
#endif
        VariableWrapper *wrapper = addLocalVar(var_decl);
        if(wrapper != NULL && current_function != NULL && var_decl->get_definition() != NULL) {
            CollectRefs *var_init = new CollectRefs(this);
            SgAssignInitializer *init = isSgAssignInitializer(var_decl->get_definition()->get_vardefn()->get_initializer());
            if(init != NULL) {
                var_init->traverse(init->get_operand());
                var_init->finish();
                Reference **reads = var_init->getReferences();
                int n_reads = var_init->getNumReferences();
                appendReads(reads, n_reads);

                Reference *new_write = new Reference(VARIABLE);
                new_write->setVar(wrapper);
                appendWrites(&new_write, 1);
            }
        }
    } else if(var_def != NULL) {
#ifdef DO_PRINT
        cout<<"Variable definition..."<<endl;
#endif
        SgInitializer *init = var_def->get_vardefn()->get_initializer();
        if(init != NULL) {
            cout<<"For init "<<init->variantT()<<endl;
            CollectRefs *var_init = new CollectRefs(this);
            var_init->traverse(init);
            var_init->finish();
            Reference **reads = var_init->getReferences();
            int n_reads = var_init->getNumReferences();
            appendWrites(reads, n_reads);
        }
    } else if(var_ref != NULL) {
#ifdef DO_PRINT
        cout<<"Variable Reference..."<<endl;
#endif
        SgVariableSymbol *sym = var_ref->get_symbol();
        SgName name = sym->get_name();

    } else if(struct_def != NULL) {
#ifdef DO_PRINT
        cout<<"Struct Definition..."<<endl;
#endif
        addStruct(struct_def);
        current_scope = struct_def;
    } else if(scope != NULL) {
#ifdef DO_PRINT
        cout<<"Scope..."<<endl;
#endif
        current_scope = scope;
    } else if(assign != NULL) {
#ifdef DO_PRINT
        cout<<"Assign..."<<endl;
        //cout<<assign->get_lhs_operand()->variantT()<<endl;
#endif

        CollectRefs *cr_write = new CollectRefs(this);
        cr_write->traverse(assign->get_lhs_operand());
        cr_write->finish();
        Reference **new_writes = cr_write->getReferences();
        int n_new_writes = cr_write->getNumReferences();
        appendWrites(new_writes, n_new_writes);

        CollectRefs *cr_read = new CollectRefs(this);
        cr_read->traverse(assign->get_rhs_operand());
        cr_read->finish();
        Reference **new_reads = cr_read->getReferences();
        int n_new_reads = cr_read->getNumReferences();
        appendReads(new_reads, n_new_reads);
    } else if(return_stmt != NULL) {
#ifdef DO_PRINT
        cout<<"Return..."<<endl;
#endif
        SgExpression *expr = return_stmt->get_expression();

        CollectRefs *cr_read = new CollectRefs(this);
        cr_read->traverse(expr);
        cr_read->finish();
        Reference **new_reads = cr_read->getReferences();
        int n_new_reads = cr_read->getNumReferences();
        appendReads(new_reads, n_new_reads);
    }

    if(current_scope != NULL) {
        scope_stack = (SgScopeStatement **)realloc(scope_stack, sizeof(SgScopeStatement *) * (scope_depth+1));
        scope_stack[scope_depth] = current_scope;
        scope_depth++;
    }
}

void analyzeTraversal::postOrderVisit(SgNode *node) {
    SgFunctionDefinition *def = isSgFunctionDefinition(node);
    SgClassDefinition *struct_def = isSgClassDefinition(node);
    if(def != NULL ) {
        if(current_function_index != -1) {
            current_function_index = -1;
            current_function = NULL;
        }
    } else if(struct_def != NULL) {
        if(current_struct_index != -1) {
            current_struct_index = -1;
            current_struct = NULL;
        }
    }
    scope_stack = (SgScopeStatement **)realloc(scope_stack, sizeof(SgScopeStatement *) * (scope_depth - 1));
    scope_depth--;
}

int analyzeTraversal::contains(Reference **arr, Reference *val, int N) {
    for(int i = 0; i < N; i++) {
        if(arr[i]->equals(val)) { 
            return 1;
        }
    }
    return 0;
}

int analyzeTraversal::duplicate(Reference **arr, Reference *find, int N) {
    for(int i = 0; i < N; i++) {
        if(arr[i]->equals(find)) return 1;
    }
    return 0;
}

void analyzeTraversal::removeDuplicates() {
    for(int i = 0; i < n_functions ;i++) {
        FunctionContext *curr_func = all_functions[i];
        for(int j = 0; j < curr_func->n_args(); j++) {
            VariableWrapper *curr_param = curr_func->ithParam(j);

            if(curr_param != NULL) {
                Reference **new_aliases = NULL;
                int num_new_aliases = 0;
                Reference **old_aliases = curr_param->getAliases();
                for(int k = 0; k < curr_param->num_aliases(); k++) {
                    if(old_aliases[k] != NULL) {
                        if(duplicate(old_aliases, old_aliases[k], k) == 0) {
                            new_aliases = (Reference **)realloc(new_aliases, sizeof(Reference *) * (num_new_aliases+1));
                            new_aliases[num_new_aliases] = old_aliases[k];
                            num_new_aliases++;
                        }
                    }
                }
                curr_param->reset_aliases(new_aliases, num_new_aliases);
            }
        }
    }
}

void analyzeTraversal::gatherAllAliases() {
    for(int i = 0; i < n_functions ;i++) {
        FunctionContext *curr_func = all_functions[i];
        for(int j = 0; j < curr_func->n_args(); j++) {
            VariableWrapper *curr_param = curr_func->ithParam(j);

            if(curr_param != NULL) {
                Reference **new_aliases = NULL;
                int num_new_aliases = 0;
                vector<Reference *> pending;
                for(int k = 0; k < curr_param->num_aliases(); k++) {
                    if(curr_param->ithAlias(k) != NULL) {
                        pending.push_back(curr_param->ithAlias(k));
                    }
                }
                while(pending.size() > 0) {
                    Reference *tmp = pending.back();
                    pending.pop_back();
                    
                    if(contains(new_aliases, tmp, num_new_aliases) == 0) {
                        new_aliases = (Reference **)realloc(new_aliases, sizeof(Reference *) * (num_new_aliases+1));
                        new_aliases[num_new_aliases] = tmp;
                        num_new_aliases++;

                        for(int k = 0; k < tmp->getVar()->num_aliases(); k++) {
                            if(tmp->getVar()->ithAlias(k) != NULL) {
                                Reference *alias_of_alias = tmp->getVar()->ithAlias(k);
                                Reference *tmp_copy = tmp->copy();
                                Reference *aoa_copy = alias_of_alias->copy();
                                tmp_copy = tmp_copy->append(aoa_copy);
                                pending.push_back(tmp_copy);
                            }
                        }
                    }
                }
                curr_param->reset_aliases(new_aliases, num_new_aliases);
            }
        }
    }
}

int main(int argc, char **argv) {
    SgProject *project = frontend(argc, argv);
    analyzeTraversal *analyze = new analyzeTraversal();
    confirmFunctions *confirm = new confirmFunctions(analyze);
    cout<<"Before confirm traversal"<<endl;
    confirm->traverseInputFiles(project);
    cout<<"After confirm traversal"<<endl;
    analyze->setFunctions(confirm->getFunctions(), confirm->getNumFunctions());
    cout<<"After setting functions "<<confirm->getNumFunctions()<<endl;
    analyze->traverseInputFiles(project);
    cout<<"After analysis traversal"<<endl;
    analyze->gatherAllAliases();
    analyze->removeDuplicates();
    cout<<"Preparing print"<<endl;
    analyze->print();

    return 0;
}
